#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-09-18


import threading
import time
import cv2
import os
import json
import argparse
import numpy as np
from typing import Union
from queue import Queue
import platform
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args, Rate
import base64
import copy
import math


class MultiImageVisNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.image_topics = self.get_param("image_topics", [
            "/vid01/sensor/image_raw",
            "/vid02/sensor/image_raw",
            "/vid03/sensor/image_raw",
            "/vid04/sensor/image_raw",
            "/vid05/sensor/image_raw",
            "/vid06/sensor/image_raw",
            "/vid07/sensor/image_raw",
            "/vid08/sensor/image_raw",
            "/vid09/sensor/image_raw",
            "/vid10/sensor/image_raw",
            "/vid11/sensor/image_raw",
            "/vid12/sensor/image_raw",
            "/vid13/sensor/image_raw",
            "/vid14/sensor/image_raw",
            "/vid15/sensor/image_raw",
            "/vid16/sensor/image_raw"
        ])
        self.show_delays = self.get_param("show_delays", False)
        self.params_help()

        self.readers = []
        self.delays = {}
        self.counts = {}
        self.images = {}
        self.images_lock = {}
        for url_ in self.image_topics:
            reader_ = Subscriber(
                url_, 'std_msgs::Null', lambda msg, url=url_: self.img_callback(msg, url),
                ip=ip, port=port
            )
            print('  ' + url_)
            self.readers.append(reader_)
            self.images[url_] = np.zeros((720, 1280, 3), dtype=np.uint8)
            self.images_lock[url_] = threading.Lock()
            self.delays[url_] = []
            self.counts[url_] = 0

        self.start()
    
    def img_callback(self, msg, url):
        img = sms2cvimg(msg)
        if 't0' in msg:
            dt = time.time() - msg['t0']
        else:
            dt = time.time() - msg['timestamp']
        with self.images_lock[url]:
            self.images[url] = img
            self.delays[url].append(dt * 1000)
            self.counts[url] += 1

    def release(self):
        BaseNode.release(self)
        for reader_ in self.readers:
            reader_.kill()

    def run(self):
        r = Rate(30)
        t = 0
        avg_delays = {}
        max_delays = {}
        hzs = {}
        while self.is_running():
            imgs = {}
            if time.time() - t >= 1.0:
                update = True
                t = time.time()
            else:
                update = False
            for url in self.image_topics:
                with self.images_lock[url]:
                    imgs[url] = self.images[url].copy()
                    if update:
                        d_arr = np.array(self.delays[url])
                        avg_delays[url] = np.mean(d_arr) if len(d_arr) > 0 else 0
                        max_delays[url] = np.max(d_arr) if len(d_arr) > 0 else 0
                        self.delays[url] = []
                        hzs[url] = self.counts[url]
                        self.counts[url] = 0

            n_imgs = len(imgs)
            l_side = int(np.ceil(np.sqrt(n_imgs)))
            sw = int(1280 / l_side)
            sh = int(720 / l_side)
            w = sw * l_side
            h = sh * l_side
            concat_img = np.zeros((h, w, 3), np.uint8)
            for i, url in enumerate(self.image_topics):
                oh, ow = imgs[url].shape[:2]
                img = cv2.resize(imgs[url], (sw, sh))
                row = i // l_side
                col = i % l_side
                concat_img[row * sh: (row+1) * sh, col * sw: (col+1) * sw, :] = img
            
            if self.show_delays:
                for i, url in enumerate(self.image_topics):
                    cv2.putText(
                        concat_img, 
                        "{}: {}x{}@{} Hz, ({:.2f} ms / {:.2f} ms) (avg / max) delay.".format(i, ow, oh, hzs[url], avg_delays[url], max_delays[url]), 
                        (10, 20 + i * 20), 
                        cv2.FONT_HERSHEY_SIMPLEX, 
                        0.5, 
                        (255, 255, 255), 
                        thickness=1, 
                        lineType=cv2.LINE_AA
                    )

            cv2.imshow("concatimg", concat_img)
            cv2.waitKey(5)
            r.sleep()

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = MultiImageVisNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()

